
# import bs4
import glob
import logging
import lxml
import math
import os

import helper


logger = logging.getLogger(__name__)


class Main():

    def __init__(self, args):
        self.args = args

    def start(self):
        class_names = []

        os.makedirs(self.args.out_dir, exist_ok=True)

        for label_fullname in glob.glob(os.path.join(self.args.in_dir, '*.xml')):
            basename = os.path.basename(label_fullname)

            parser = helper.LabelXml(class_names, label_fullname)
            _, img_h, img_w = parser.get_chw()
            img_hw = (img_h, img_w)

            boxes = []
            for class_id, cxywh in parser.get_boxes(img_hw):
                nxywh = helper.normalize_xywh(cxywh, img_hw)
                boxes.append((class_id, nxywh))

            if not boxes:
                continue

            out_fullname = os.path.join(self.args.out_dir, basename.replace('.xml', '.txt'))
            with open(out_fullname, 'w') as f:
                for class_id, nxywh in boxes:
                    f.write('%d %f %f %f %f\n' % (class_id, *nxywh))

        if self.args.classes_file is not None:
            classes_fullname = self.args.classes_file
        else:
            classes_fullname = os.path.join(self.args.out_dir, 'classes.txt')
        with open(classes_fullname, 'w') as f:
            for class_name in class_names:
                f.write('%s\n' % class_name)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--in-dir', type=str, default=r'C:\my\ws\static\fair1m1.0\train-part1\labelXml')
    parser.add_argument('--out-dir', type=str, default=r'C:\my\ws\static\fair1m1.0\train-part1\labelXml-txt')
    parser.add_argument('--classes-file', type=str, default=None)
    args = parser.parse_args()
    Main(args).start()
